#include "SocketReaderWriter.h"

#include <boost/bind.hpp>

namespace AlgoTrading
{
	namespace ConnectionLibrary
	{

		SocketReaderWriter::SocketReaderWriter(boost::asio::io_service& io_service, boost::asio::ip::tcp::socket& socket, SocketNotificationHandler& notificationHandler)
			: io_service(io_service)
			, socket(socket)
			, disconnected(false)
			, writeInProgress(false)
			, unsentBuffersCount(0)
			, sentBuffersCount(0)
			, streamBuf(0)
			, notificationHandler(notificationHandler)
			, strand(new boost::asio::strand(io_service))
		{
		}

		SocketReaderWriter::~SocketReaderWriter(void)
		{
		}

		void SocketReaderWriter::Disconnect()
		{
			strand->dispatch(boost::bind(&SocketReaderWriter::DisconnectImpl, shared_from_this()));
		}

		void SocketReaderWriter::DisconnectImpl()
		{
			if (!disconnected)
			{
				socket.shutdown(boost::asio::ip::tcp::socket::shutdown_both);
				socket.close();
				notificationHandler.OnDisconnect();
				disconnected = true;
			}
		}

		void SocketReaderWriter::SendBuffer(OutputMessageBufferHolderPtr outputMessageBufferHolder)
		{	
			strand->post(boost::bind(&SocketReaderWriter::PerformSendRequest, shared_from_this(), outputMessageBufferHolder));
		}

		void SocketReaderWriter::StartReading()
		{
			strand->post(boost::bind(&SocketReaderWriter::PerformRead, shared_from_this()));
		}

		void SocketReaderWriter::PerformSendRequest(OutputMessageBufferHolderPtr outputMessageBufferHolder)
		{
			cacheForWriting.push_back(outputMessageBufferHolder);
			++unsentBuffersCount;
			if (!writeInProgress)
			{
				ActiveWrite();
			}
		}

		void SocketReaderWriter::ActiveWrite()
		{
			std::vector<boost::asio::const_buffer> buffersToWrite;
			buffersToWrite.reserve(unsentBuffersCount * 2);
			std::deque<OutputMessageBufferHolderPtr>::iterator it = cacheForWriting.begin();
			for(int i = 0; i < unsentBuffersCount; ++i)
			{
				OutputMessageBufferHolderPtr buffersHolder = *it;

				buffersToWrite.push_back(boost::asio::buffer(buffersHolder->GetMessageLength(), sizeof(::google::protobuf::uint32)));
				buffersToWrite.push_back(buffersHolder->GetMessageBuffer());

				++it;
			}
			sentBuffersCount = unsentBuffersCount;
			unsentBuffersCount = 0;

			boost::asio::async_write(socket, buffersToWrite,
				make_custom_alloc_handler(allocator,
				strand->wrap(
				boost::bind(&SocketReaderWriter::HandleWrite,
				shared_from_this(),
				boost::asio::placeholders::error,
				boost::asio::placeholders::bytes_transferred))));

			writeInProgress = true;
		}

		void SocketReaderWriter::HandleWrite(const boost::system::error_code& error,size_t bytes_transferred)
		{
			if (error)
			{
				if (error != boost::asio::error::operation_aborted)
				{
					notificationHandler.OnError(error);
					DisconnectImpl();
				}

				cacheForWriting.clear();
				sentBuffersCount = 0;
				unsentBuffersCount = 0;

				return;
			}
			else
			{
				for(int i = 0; i < sentBuffersCount; ++i)
				{
					cacheForWriting.pop_front();
				}

				sentBuffersCount = 0;
			}

			writeInProgress = false;

			if (unsentBuffersCount > 0)
			{
				ActiveWrite();
			}
		}

		void SocketReaderWriter::HandleReadDataLength(const boost::system::error_code& error)
		{
			if (error)
			{
				if (error != boost::asio::error::operation_aborted)
				{
					notificationHandler.OnError(error);
					DisconnectImpl();
				}
			}
			else
			{
				streamBuf.reset(new boost::asio::streambuf());
				boost::asio::async_read(socket, 
					streamBuf->prepare((size_t)messageLength),
					make_custom_alloc_handler(allocator,
					strand->wrap(
					boost::bind(&SocketReaderWriter::HandleReadMessage,
					shared_from_this(),
					boost::asio::placeholders::error))));
			}
		}

		void SocketReaderWriter::HandleReadMessage(const boost::system::error_code& error)
		{
			if (error)
			{
				if (error != boost::asio::error::operation_aborted)
				{
					notificationHandler.OnError(error);
					DisconnectImpl();
				}
			}
			else
			{
				std::istream stream(streamBuf.get());
				streamBuf->commit(messageLength);
				notificationHandler.ParseMessageData(stream);
				stream.clear();
				streamBuf.reset();
				PerformRead();
			}
		}

		void SocketReaderWriter::PerformRead()
		{
			boost::asio::async_read(socket,
				boost::asio::buffer((void*) &messageLength,
				sizeof(::google::protobuf::uint32)),
				make_custom_alloc_handler(allocator,
				strand->wrap(
				boost::bind(&SocketReaderWriter::HandleReadDataLength,
				shared_from_this(),
				boost::asio::placeholders::error))));
		}

	}
}
